import torch
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import gzip
import csv
import time
from torch.nn.utils.rnn import pack_padded_sequence
import math
#可不加
import os
os.environ["KMP_DUPLICATE_LIB_OK"]  =  "TRUE"

HIDDEN_SIZE = 100
BATCH_SIZE = 256
#GRU层数为2
N_LAYER = 2
#训练100轮
N_EPOCHS = 100
#字符长度128
N_CHARS = 128
USE_GPU = True

class RNNClassifier(torch.nn.Module):
    #bidirectional 是否为双向神经网络
    def __init__(self, input_size, hidden_size, output_size, n_layers =1 , bidirectional = True):
        super(RNNClassifier, self).__init__()
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.n_directions = 2 if bidirectional else 1
        #Embedding层输入 （SeqLen，BatchSize）
        #Embedding层输出 （SeqLen，BatchSize，HiddenSize）
        #将原先样本总数为SeqLen，批量数为BatchSize的数据，转换为HiddenSize维的向量
        self.embedding = torch.nn.Embedding(input_size, hidden_size)
        #bidirection用于表示神经网络是单向还是双向
        self.gru = torch.nn.GRU(hidden_size, hidden_size, n_layers, bidirectional = bidirectional)
        #线性层需要*direction 双向最后输出要拼接所以要乘2
        self.fc = torch.nn.Linear(hidden_size * self.n_directions, output_size)

    def _init_hidden(self):
        #创建全0初始隐层
        hidden = torch.zeros(self.n_layers * self.n_directions, BATCH_SIZE, self.hidden_size)

        #return create_tensors(hidden)

    def forward(self, input, seq_length):
        #对input进行转置
        input = input.t()
        batch_size = input.size(1)

        #（n_Layer * nDirections, BatchSize, HiddenSize）
        hidden = self._init_hidden(batch_size)
        #(SeqLen, BatchSize, HiddenSize)
        embedding = self.embedding(input)

        #对数据计算过程提速
        #需要得到嵌入层的结果（输入数据）及每条输入数据的长度
        gru_input = pack_padded_sequence(embedding, seq_length)

        output, hidden = self.gru(gru_input, hidden)

        #如果是双向神经网络会有h_N^f以及h_1^b两个hidden
        if self.n_directions == 2:
            hidden_cat = torch.cat([hidden[-1], hidden[-2]], dim=1)
        else:
            hidden_cat = hidden[-1]

        fc_output = self.fc(hidden_cat)

        return fc_output

class NameDataset(Dataset):
    def __init__(self, is_train_set=True):

        #读数据 判断是不是训练集
        filename = 'names_train.csv.gz' if is_train_set else 'names_test.csv.gz'
        with gzip.open(filename, 'rt') as f:
            reader = csv.reader(f)
            rows = list(reader)

        #rows中每一个元素都是数据元组（name,country）,将其中的name和country提取出来，并记录数量
        self.names = [row[0] for row in rows]
        self. len = len(self.names)
        self.countries = [row[1] for row in rows]

        #将country转换成索引
        #列表->集合->排序->列表->字典
        self.country_list = list(sorted(set(self.countries)))
        #getCountryDict把列表变为词典
        self.country_dict = self.getCountryDict()
        #获取长度
        self.country_num = len(self.country_list)

    #获取键值对，country(key)-index(value)
    def __getitem__(self, index):
        #输出一个名字字符串和国家索引
        return self.names[index], self.country_dict[self.countries[index]]

    def __len__(self):
        return self.len

    def getCountryDict(self):
        #新建一个字典
        country_dict = dict()
        for idx,country_name in enumerate(self.country_list, 0):
            country_dict[country_name]=idx
        #输出国家字典
        return country_dict

    #根据索引返回国家名
    def idx2country(self, index):
        return self.country_list[index]

    #返回国家数目
    def getCountriesNum(self):
        return self.country_num

trainset = NameDataset(is_train_set = True)
trainloader = DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True)

testset = NameDataset(is_train_set=False)
testloader = DataLoader(testset, batch_size=BATCH_SIZE, shuffle=False)

#最终的输出维度 获得国家个数
N_COUNTRY = trainset.getCountriesNum()


def time_since(since):
    s = time.time() - since
    m = math.floor(s / 60)
    s -= m*60
    return '%dm %ds' % (m, s)

if __name__ == '__main__':
    '''
    N_CHARS：字符数量，英文字母转变为One-Hot向量
    HIDDEN_SIZE：GRU输出的隐层的维度 GRU是RNN的升级版
    N_COUNTRY：分类的类别总数
    N_LAYER：GRU层数
    '''
    classifier = RNNClassifier(N_CHARS, HIDDEN_SIZE, N_COUNTRY, N_LAYER)
    #迁移至GPU
    if USE_GPU:
        device = torch.device("cuda:0")
        classifier.to(device)

    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001)

    #开始时间
    start = time.time()
    print("Training for %d epochs ... " % N_EPOCHS)
    #记录训练准确率
    acc_list = []
    for epoch in range(1, N_EPOCHS+1):
        print("")
        # #训练模型
        # trainModel()
        # #检测模型
        # acc = testModel()
        # acc_list.append(acc)

    #绘制图像
    epoch = np.arange(1, len(acc_list)+1, 1)
    acc_list = np.array(acc_list)
    plt.plot(epoch, acc_list)
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.grid()
    plt.show()